from a2a.server.request_handlers import DefaultRequestHandler
from a2a.types import AgentCard, AgentCapabilities, AgentSkill
from a2a.server.tasks import BasePushNotificationSender, InMemoryPushNotificationConfigStore, InMemoryTaskStore
from a2a.server.apps import A2AStarletteApplication
from agent_executor import CurrencyAgentExecutor
from travel import TravelAgent
from utils import Model, ModelType
from pydantic import BaseModel
from dotenv import load_dotenv
import httpx
import uvicorn
import torch
import os

load_dotenv()
model_name = os.getenv("DEEPSEEK_NAME")
model_key = os.getenv("DEEPSEEK_KEY")
model_url = os.getenv("DEEPSEEK_URL")
skill = AgentSkill(
    id='convert_currency',
    name='Currency Exchange Rates Tool',
    description='Helps with exchange values between various currencies',
    tags=['currency conversion', 'currency exchange'],
    examples=['What is exchange rate between USD and GBP?'],
)

agent_card = AgentCard(
    name="高德地图工具 Agent",
    description="查询地址与位置信息",
    url="http://0.0.0.0:8001/",
    version="1.0.0",
    capabilities=AgentCapabilities(streaming=False),
    skills=[skill],
    default_input_modes=['text', 'text/plain'],
    default_output_modes=['text', 'text/plain'],
)


# 结构化输出
class WeatherResponse(BaseModel):
    conditions: str


httpx_client = httpx.AsyncClient()

push_config_store = InMemoryPushNotificationConfigStore()
push_sender = BasePushNotificationSender(httpx_client=httpx_client,
                                         config_store=push_config_store)

model = Model(ModelType.HTTP, torch.float32,model_name, model_url=model_url, model_key=model_key)
agent = TravelAgent(model, [], "", WeatherResponse)

request_handler = DefaultRequestHandler(
    agent_executor=CurrencyAgentExecutor(agent),  # Agent执行器
    task_store=InMemoryTaskStore(),  # 内存中的任务存储
    push_config_store=push_config_store,
    push_sender=push_sender
)

server = A2AStarletteApplication(
    agent_card=agent_card,
    http_handler=request_handler,
    extended_agent_card=agent_card,
)

if __name__ == "__main__":
    uvicorn.run(server.build(), host="0.0.0.0", port=30001)
